import torch
from torch import nn


class FingerPrintModule_old(nn.Module):
    """分子指纹模块，用于分别提取MACCS和ECFP特征"""

    def __init__(self):
        super(FingerPrintModule_old, self).__init__()

        # MACCS指纹处理通道（167 -> 256*2）
        self.maccs_net = nn.Sequential(
            nn.Linear(167, 167 * 2),
            nn.ReLU(),
            nn.Linear(167 * 2, 256)
        )

        # ECPF指纹处理通道（1024 -> 256）
        self.ecfp_net = nn.Sequential(
            nn.Linear(2048, 1024),
            nn.ReLU(),
            nn.Linear(1024, 512),
            nn.ReLU(),
            nn.Linear(512, 256)
        )

        # 注意力机制融合
        self.gate = nn.Sequential(
            nn.Linear(256 * 2, 1),  # 输出单个权重
            nn.Sigmoid()
        )

    def forward(self, maccs_feat, ecfp_feat):
        # 处理MACCS指纹
        maccs_feat = self.maccs_net(maccs_feat)

        # 处理ECFP指纹
        ecfp_feat = self.ecfp_net(ecfp_feat)

        # 拼接特征
        combined = torch.cat([maccs_feat, ecfp_feat], dim=1)
        gate_weight = self.gate(combined)
        fused_feat = gate_weight * maccs_feat + (1 - gate_weight) * ecfp_feat

        return fused_feat

class FingerPrintModule(nn.Module):
    """分子指纹模块，用于分别提取MACCS和ECFP特征"""
    
    def __init__(self, 
                 maccs_input_dim=167,
                 maccs_hidden_dim=334,  # 167*2
                 maccs_output_dim=256,
                 ecfp_input_dim=2048,
                 ecfp_hidden_dims=(1024, 512),
                 ecfp_output_dim=256,
                 use_attention=True):
        super(FingerPrintModule, self).__init__()
        
        # MACCS处理通道
        self.maccs_net = nn.Sequential(
            nn.Linear(maccs_input_dim, maccs_hidden_dim),
            nn.ReLU(),
            nn.Linear(maccs_hidden_dim, maccs_output_dim)
        )

        # ECFP处理通道
        ecfp_layers = []
        input_dim = ecfp_input_dim
        for hidden_dim in ecfp_hidden_dims:
            ecfp_layers.extend([
                nn.Linear(input_dim, hidden_dim),
                nn.ReLU()
            ])
            input_dim = hidden_dim
        ecfp_layers.append(nn.Linear(input_dim, ecfp_output_dim))
        self.ecfp_net = nn.Sequential(*ecfp_layers)

        # 注意力机制
        self.use_attention = use_attention
        if self.use_attention:
            self.gate = nn.Sequential(
                nn.Linear(maccs_output_dim + ecfp_output_dim, 1),
                nn.Sigmoid()
            )

    def forward(self, maccs_feat, ecfp_feat):
        maccs_feat = self.maccs_net(maccs_feat)
        ecfp_feat = self.ecfp_net(ecfp_feat)

        if self.use_attention:
            combined = torch.cat([maccs_feat, ecfp_feat], dim=1)
            gate_weight = self.gate(combined)
            return gate_weight * maccs_feat + (1 - gate_weight) * ecfp_feat
        return (maccs_feat + ecfp_feat) / 2  # 不使用注意力时取平均
    
    
if __name__ == "__main__":
    # 设置随机种子以保证可重复性
    torch.manual_seed(42)

    # 创建测试数据
    batch_size = 16
    maccs_data = torch.randn(batch_size, 167)  # MACCS指纹数据
    ecfp_data = torch.randn(batch_size, 2048)  # ECFP指纹数据

    # 初始化模型
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = FingerPrintModule(device).to(device)

    output = model(maccs_data, ecfp_data)
